import os 
import shelve
import time
from tqdm import tqdm

def sampleGraphStruct(root_path,lable_path,save_path):
    """
        获得样本的图结构分布
    """
    a=time.time()
    samples={}
    fam_lables={}
    numberCount = []
    for i in os.listdir(root_path):
        tmp=i.split(".")
        if tmp[1]=="dir":
            samples[tmp[0]]=1
    
    with open(lable_path,"r") as file:
        for line in tqdm(file):
            tmp=line.split("\t")
            sampleName = tmp[0]
            sampleLable = tmp[1]
            if sampleName in samples:
                with shelve.open(root_path+"//"+sampleName) as file:
                    cg=file["cg"]
                    funcType=file["funcType"]
                
                a,b,c = 0, 0, 0
                for ind in funcType:
                    if funcType[ind] == "local":
                        a+=1
                    elif funcType[ind] == "dynamic import":
                        c+=1
                    else:
                        b+=1
                numberCount.append((len(cg),a,b,c))
    
                print("已处理样本{},长度{},分布：{}-{}-{}".format(sampleName,len(cg),a,b,c))

                if sampleLable not in fam_lables:
                    fam_lables[sampleLable]=[(sampleName,len(cg),a,b,c)]
                else:
                    fam_lables[sampleLable].append((sampleName,len(cg),a,b,c))
    # 计算数据分布
    res = [[0,0,0,0,0],[0,0,0,0,0],[0,0,0,0,0],[0,0,0,0,0]]
    for number,a,b,c in numberCount:
        ind  = 3
        if number < 100:
            ind = 0
        elif number <1000:
            ind = 1
        elif number < 10000:
            ind = 2
        res[ind][0]+=1   
        res[ind][1]+=number    
        res[ind][2]+=a    
        res[ind][3]+=b    
        res[ind][4]+=c

    b=time.time()
    with  open(save_path,"a") as file:
        file.write("生成耗时{}\n".format(b-a))
        file.write("所有样本的数量：{}\n".format(len(samples)))
        file.write("家族的数量：{}\n".format(len(fam_lables)))
        file.write("结点统计信息：\n")
        file.write("0~100:\t{}\n".format(res[0]))
        file.write("100~1000:\t{}\n".format(res[1]))
        file.write("1000~10000:\t{}\n".format(res[2]))
        file.write("10000以上:\t{}\n".format(res[3]))

        for i in fam_lables.keys():
            file.write(i+"\t数量: "+str(len(fam_lables[i]))+"\n")
            for name,lth,a,b,c in fam_lables[i]:
                file.write("\t"+name+" 分布： "+str(lth)+" "+str(a)+" "+str(b)+" "+str(c))
            file.write("\n")
            
if __name__=="__main__":
    save_path = r"/home/cyw/projects/function_sim_project/all_data/graphStruct.txt"
    # sampleGraphStruct(r"/mnt/mydisk1/chenyongwei/sampleDatas/functionSim",r"/home/cyw/projects/function_sim_project/all_data/fam_lables.txt")
    sampleGraphStruct(r"/mnt/mydisk1/chenyongwei/malware/BODMAS_dataset/HGMSim",r"/home/cyw/projects/function_sim_project/all_data/bodmasMD5FamLable.txt",save_path)